#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EndoSight-UC AI推理服务器 - 优化版本
使用训练好的UCEIS评分模型进行内镜图像分析
"""

import os
import sys
import logging
import traceback
import time
from pathlib import Path
from typing import Dict, Any, Optional
import uuid
import json

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
from torchvision.models import ResNet50_Weights
from transformers import BertModel, BertTokenizer
from PIL import Image
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn
from pydantic import BaseModel

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('ai_inference_optimized.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 全局变量
model = None
device = None
tokenizer = None

# UCEIS评分映射 (0-8分)
SEVERITY_MAPPING = {
    0: "缓解期",
    1: "轻度",
    2: "轻度",
    3: "中度",
    4: "中度",
    5: "重度",
    6: "重度",
    7: "重度",
    8: "重度"
}

# 响应模型
class PredictionResponse(BaseModel):
    task_id: str
    uceis_score: int
    severity: str
    confidence: float
    detailed_scores: Dict[str, Any]
    processing_time: float
    success: bool
    message: str

class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    device: str
    message: str

# 导入训练时使用的模型结构
sys.path.insert(0, r"D:\VSCodeProjects\UC_System\uc_model")
try:
    from models.fusion_model import EndoMultimodalModel
    logger.info("成功导入训练时使用的EndoMultimodalModel")
except ImportError as e:
    logger.error(f"无法导入训练模型: {e}")
    sys.exit(1)

def load_model() -> bool:
    """加载训练好的模型"""
    global model, device, tokenizer

    try:
        logger.info("开始加载UCEIS模型...")

        # 设备选择
        if torch.cuda.is_available():
            device = torch.device('cuda')
            logger.info(f"使用GPU: {torch.cuda.get_device_name(0)}")
        else:
            device = torch.device('cpu')
            logger.info("使用CPU进行推理")

        # 初始化tokenizer
        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        logger.info("BERT tokenizer加载成功")

        # 创建模型实例 - 使用与训练时相同的参数
        logger.info("初始化模型结构...")
        model = EndoMultimodalModel(
            image_pretrained=True,
            text_pretrained='bert-base-uncased',
            feature_dim=256,
            num_classes=8
        )
        logger.info("模型结构创建成功")

        # 加载训练好的权重
        model_path = Path(r"D:\VSCodeProjects\UC_System\uc_model\final_weights\uceis_best_model.pth")

        if not model_path.exists():
            logger.error(f"模型文件不存在: {model_path}")
            return False

        logger.info(f"加载模型权重: {model_path}")
        checkpoint = torch.load(model_path, map_location=device)

        # 加载模型状态字典
        if 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'], strict=True)
            logger.info("成功加载model_state_dict")
        else:
            model.load_state_dict(checkpoint, strict=True)
            logger.info("直接加载checkpoint权重")

        # 移动到设备并设置为评估模式
        model.to(device)
        model.eval()

        # 记录模型信息
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

        logger.info(f"模型加载成功！")
        logger.info(f"总参数量: {total_params:,}")
        logger.info(f"可训练参数: {trainable_params:,}")
        logger.info(f"训练轮数: {checkpoint.get('epoch', '未知')}")
        logger.info(f"最佳F1分数: {checkpoint.get('best_val_f1', '未知'):.4f}")

        return True

    except Exception as e:
        logger.error(f"模型加载失败: {str(e)}")
        logger.error(f"错误详情: {traceback.format_exc()}")
        return False

def preprocess_image(image_file) -> torch.Tensor:
    """图像预处理"""
    try:
        # 读取图像
        image = Image.open(image_file.file).convert('RGB')
        logger.info(f"成功加载图像: {image_file.filename}, 尺寸: {image.size}")

        # 预处理变换
        transform = transforms.Compose([
            transforms.Resize((512, 512)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                               std=[0.229, 0.224, 0.225])
        ])

        # 应用变换
        image_tensor = transform(image).unsqueeze(0)  # 添加batch维度
        return image_tensor.to(device)

    except Exception as e:
        logger.error(f"图像预处理失败: {str(e)}")
        raise HTTPException(status_code=400, detail=f"图像预处理失败: {str(e)}")

def predict_uceis_score(image_tensor: torch.Tensor) -> Dict[str, Any]:
    """预测UCEIS评分"""
    try:
        start_time = time.time()

        with torch.no_grad():
            # 准备文本输入（使用默认的医学描述）
            default_text = "endoscopic image of colon mucosa for UCEIS scoring"
            text_inputs = tokenizer(
                default_text,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=128
            )

            # 移动到正确设备
            for k, v in text_inputs.items():
                text_inputs[k] = v.to(device)

            # 模型推理 - 注意模型的输入格式
            outputs = model(image_tensor, text_inputs)  # 返回 (B, 8) 的logits

            # 获取预测结果
            logits = outputs.squeeze(0)  # 移除batch维度，得到 (8,)
            probabilities = F.softmax(logits, dim=0)

            # 获取预测分数（0-7对应0-7分，但我们需要映射到0-8分）
            predicted_class = torch.argmax(probabilities).item()

            # 将0-7的类别映射到0-8的UCEIS分数
            # 注意：训练时使用的标签可能是0-7，需要映射到0-8
            uceis_score = predicted_class  # 如果训练是0-7对应0-7分
            # 如果需要映射到0-8，可以使用：uceis_score = min(predicted_class, 8)

            # 计算置信度
            confidence = probabilities[predicted_class].item()

            # 构建详细分数（由于模型只输出总分，其他子评分为估算）
            detailed_scores = {
                'total': {
                    'score': uceis_score,
                    'confidence': confidence,
                    'probabilities': probabilities.cpu().numpy().tolist()
                },
                'vascular': {
                    'score': min(uceis_score, 2),  # 简化的子评分估算
                    'confidence': confidence * 0.8,
                    'note': '基于总分估算'
                },
                'bleeding': {
                    'score': min(max(uceis_score - 2, 0), 2),
                    'confidence': confidence * 0.8,
                    'note': '基于总分估算'
                },
                'erosion': {
                    'score': min(max(uceis_score - 4, 0), 3),
                    'confidence': confidence * 0.8,
                    'note': '基于总分估算'
                }
            }

            # 获取严重程度
            severity = SEVERITY_MAPPING.get(uceis_score, "未知")

            processing_time = time.time() - start_time

            return {
                'uceis_score': uceis_score,
                'severity': severity,
                'confidence': confidence,
                'detailed_scores': detailed_scores,
                'processing_time': processing_time
            }

    except Exception as e:
        logger.error(f"预测失败: {str(e)}")
        logger.error(f"错误详情: {traceback.format_exc()}")
        raise HTTPException(status_code=500, detail=f"预测失败: {str(e)}")

# 创建FastAPI应用
app = FastAPI(
    title="EndoSight-UC AI推理服务器",
    description="溃疡性结肠炎内镜严重程度指数(UCEIS)评分AI服务",
    version="2.0.0"
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 性能优化：设置请求大小限制
from fastapi import Request
import time

@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    return response

@app.on_event("startup")
async def startup_event():
    """启动时加载模型"""
    success = load_model()
    if success:
        logger.info("✅ 推理服务器启动成功！")
    else:
        logger.error("❌ 模型加载失败，服务器将无法正常工作")

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """健康检查"""
    return HealthResponse(
        status="healthy" if model is not None else "unhealthy",
        model_loaded=model is not None,
        device=str(device) if device else "unknown",
        message="模型已加载" if model else "模型未加载"
    )

@app.post("/v1/predict", response_model=PredictionResponse)
async def predict_uceis(file: UploadFile = File(...)):
    """预测UCEIS评分"""
    if model is None:
        raise HTTPException(status_code=503, detail="模型未加载，请检查服务器状态")

    # 验证文件类型
    if not file.content_type or not file.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="请上传图像文件")

    # 生成任务ID
    task_id = str(uuid.uuid4())

    try:
        # 预处理图像
        image_tensor = preprocess_image(file)

        # 预测
        result = predict_uceis_score(image_tensor)

        # 构建响应
        response = PredictionResponse(
            task_id=task_id,
            uceis_score=result['uceis_score'],
            severity=result['severity'],
            confidence=result['confidence'],
            detailed_scores=result['detailed_scores'],
            processing_time=result['processing_time'],
            success=True,
            message="预测成功"
        )

        logger.info(f"任务 {task_id} 预测完成: UCEIS={result['uceis_score']}, 置信度={result['confidence']:.3f}")
        return response

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"任务 {task_id} 处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")

@app.get("/")
async def root():
    """根路径"""
    return {
        "service": "EndoSight-UC AI推理服务器",
        "status": "running",
        "version": "2.0.0",
        "model_loaded": model is not None,
        "endpoints": {
            "/health": "健康检查",
            "/v1/predict": "UCEIS评分预测",
            "/docs": "API文档"
        }
    }

if __name__ == "__main__":
    logger.info("启动EndoSight-UC AI推理服务器...")
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=5001,
        reload=False,
        log_level="info",
        workers=2  # 增加工作进程数以提高并发处理能力
    )